/*
 * hmm_test.cc
 *
 *  Created on: Sep 8, 2011
 *      Author: tqlong
 */

#include <iostream>
#include <boost/program_options.hpp>
#include <boost/date_time/posix_time/posix_time.hpp>

#include "../dhmm.h"
#include "../dur_hmm.h"
#include "../dur_hmm_markov.h"
#include "../position_model.h"
#include "../anneal.h"

namespace po = boost::program_options;
po::options_description desc("Allowed options");
po::variables_map vm;

void process_options(int argc, char** argv);
void usage(int exitcode = 0)
{
  std::cout << desc << std::endl;
  exit(exitcode);
}

void classic_hmm_test()
{
  std::cout << " ================ CLASSIC HMM ================ " << std::endl;
  typedef DiscreteHMM hmm_t;
  typedef DiscreteHMM::joint_seq_t seq_t;
  typedef DiscreteHMM::vseq_t vseq_t;
  hmm_t hmm;
  hmm.load(vm["hmmfile"].as<std::string>().c_str());
  hmm.print_params();

  seq_t seq;
  hmm.generate(50, seq);
  std::cout << "Generate sequence: " << std::endl;
  seq.print();

  double prob = hmm.seq_probability(seq.oseq);
  std::cout << "prob = " << Log(prob) << std::endl;

  prob = hmm.seq_probability_reverse(seq.oseq);
  std::cout << "prob(reverse) = " << Log(prob) << std::endl;

  seq_t map_seq(seq.oseq);
  hmm.map(map_seq);
  std::cout << "MAP sequence: " << std::endl;
  map_seq.print();

  int N = 100;
  int T = 500;
  vseq_t data;
  for (int n = 0; n < N; n++)
  {
    seq_t seq;
    hmm.generate(T, seq);
    data.push_back(seq.oseq);
  }

  hmm_t trained_hmm = hmm;
  srand(time(NULL));
  trained_hmm.rand();
  std::cout << "Before: " << std::endl;
  trained_hmm.print_params();

  trained_hmm.train(data);
  std::cout << "Trained HMM: " << std::endl;
  trained_hmm.print_params();
}

void duration_hmm_test()
{
  std::cout << " ================ Duration HMM ================ " << std::endl;
  typedef DurationHMM hmm_t;
  typedef hmm_t::joint_seq_t seq_t;
  typedef hmm_t::vseq_t vseq_t;
  hmm_t hmm;
  hmm.load(vm["hmmfile"].as<std::string>().c_str());
  hmm.print_params();
  hmm.save(vm["output"].as<std::string>().c_str());

  seq_t seq;
  hmm.generate(10, seq);
  std::cout << "Generate sequence: " << std::endl;
  seq.print();

  double prob = hmm.seq_probability(seq.oseq);
  std::cout << "prob = " << Log(prob) << std::endl;

  prob = hmm.seq_probability_reverse(seq.oseq);
  std::cout << "prob(reverse) = " << Log(prob) << std::endl;

  seq_t map_seq(seq.oseq);
  hmm.map(map_seq);
  std::cout << "MAP sequence: " << std::endl;
  map_seq.print();

  int N = 100;
  int T = 200;
  vseq_t data;
  for (int n = 0; n < N; n++)
  {
    seq_t seq;
    hmm.generate(T, seq);
    data.push_back(seq.oseq);
  }

  hmm_t trained_hmm = hmm;
  srand(time(NULL));
  trained_hmm.rand();
  std::cout << "Before: " << std::endl;
  trained_hmm.print_params();

  trained_hmm.set_prefix(vm["prefix"].as<std::string>());
  trained_hmm.train(data);
  std::cout << "Trained HMM: " << std::endl;
  trained_hmm.print_params();
}

void duration_hmm_compare_test()
{
  std::cout << " ================ Duration HMM convergence ================ " << std::endl;
  typedef DurationHMM hmm_t;

  hmm_t hmm;
  hmm.load(vm["output"].as<std::string>().c_str());
  for (int iter = 0; iter < 40; iter++)
  {
    std::stringstream s;
    s << vm["prefix"].as<std::string>().c_str() << iter << ".txt";

    hmm_t iter_hmm;
    iter_hmm.load(s.str().c_str());
    double diff = hmm.diff(iter_hmm);
    std::cout << "iter=" << iter << " diff=" << diff << std::endl;

    std::stringstream sf;
    sf << vm["prefix"].as<std::string>().c_str() << "diff.txt";
    std::ofstream llfile(sf.str().c_str(), std::ios::app);
    llfile << iter << " " << diff << std::endl;
    llfile.close();
  }
}

void duration_hmm_markov_test()
{
  std::cout << " ================ Duration HMM ================ " << std::endl;
  typedef DurationHMMMarkov hmm_t;
  typedef hmm_t::joint_seq_t seq_t;
  typedef hmm_t::vseq_t vseq_t;
  hmm_t hmm;
  hmm.load(vm["hmmfile"].as<std::string>().c_str());
  hmm.print_params();
  hmm.save(vm["output"].as<std::string>().c_str());

  seq_t seq;
  hmm.generate(10, seq);
  std::cout << "Generate sequence: " << std::endl;
  seq.print();

  double prob = hmm.seq_probability(seq.oseq);
  std::cout << "prob = " << Log(prob) << std::endl;

  prob = hmm.seq_probability_reverse(seq.oseq);
  std::cout << "prob(reverse) = " << Log(prob) << std::endl;

  seq_t map_seq(seq.oseq);
  hmm.map(map_seq);
  std::cout << "MAP sequence: " << std::endl;
  map_seq.print();

  int N = 100;
  int T = 200;
  vseq_t data;
  for (int n = 0; n < N; n++)
  {
    seq_t seq;
    hmm.generate(T, seq);
    data.push_back(seq.oseq);
  }

  hmm_t trained_hmm = hmm;
  srand(time(NULL));
  trained_hmm.rand();
  std::cout << "Before: " << std::endl;
  trained_hmm.print_params();

  trained_hmm.set_prefix(vm["prefix"].as<std::string>());
  trained_hmm.train(data);
  std::cout << "Trained HMM: " << std::endl;
  trained_hmm.print_params();
}

void duration_hmm_markov_compare_test()
{
  std::cout << " ================ Duration HMM convergence ================ " << std::endl;
  typedef DurationHMMMarkov hmm_t;

  hmm_t hmm;
  hmm.load(vm["output"].as<std::string>().c_str());
  for (int iter = 0; iter < 50; iter++)
  {
    std::stringstream s;
    s << vm["prefix"].as<std::string>().c_str() << iter << ".txt";

    hmm_t iter_hmm;
    iter_hmm.load(s.str().c_str());
    double diff = hmm.diff(iter_hmm);
    std::cout << "iter=" << iter << " diff=" << diff << std::endl;

    std::stringstream sf;
    sf << vm["prefix"].as<std::string>().c_str() << "diff.txt";
    std::ofstream llfile(sf.str().c_str(), std::ios::app);
    llfile << iter << " " << diff << std::endl;
    llfile.close();
  }
}

void motif_test()
{
  std::cout << " ================ Motif with Duration HMM ================ " << std::endl;
  typedef DurationHMM hmm_t;
  typedef hmm_t::joint_seq_t seq_t;
  typedef hmm_t::vseq_t vseq_t;
  hmm_t hmm;
  hmm.load(vm["hmmfile"].as<std::string>().c_str());
  hmm.print_params();

  seq_t seq;
  hmm.generate(8, seq);
  std::cout << "Generate sequence: " << std::endl;
  seq.print();

  double prob = hmm.seq_probability(seq.oseq);
  std::cout << "prob = " << Log(prob) << std::endl;

  prob = hmm.seq_probability_reverse(seq.oseq);
  std::cout << "prob(reverse) = " << Log(prob) << std::endl;

  seq_t map_seq(seq.oseq);
  hmm.map(map_seq);
  std::cout << "MAP sequence: " << std::endl;
  map_seq.print();

  int N = 1000;
  int T = 8;
  vseq_t data;
  std::ofstream output(vm["seqfile"].as<std::string>().c_str());
  const char *atcg = "ATCG";
  for (int n = 0; n < N; n++)
  {
    seq_t seq;
    while (seq.oseq.size() < 11) 
      hmm.generate(T, seq);
    data.push_back(seq.oseq);
    output << ">" << n+1  << std::endl;
    for (unsigned int i = 0; i < seq.oseq.size()-1; i++) {
      output << atcg[seq.oseq[i]];
    }
    output << std::endl;
  }

  hmm_t trained_hmm = hmm;
  srand(time(NULL));
  trained_hmm.rand();
  std::cout << "Before: " << std::endl;
  trained_hmm.trans_ = hmm.trans_; // fix the transmission probability table
  trained_hmm.start_ = hmm.start_; // and the starting probabilities
  trained_hmm.print_params();

  trained_hmm.set_prefix(vm["prefix"].as<std::string>());
  trained_hmm.train(data,  vm["iter"].as<int>());
  std::cout << "Trained HMM: " << std::endl;
  trained_hmm.print_params();

  {
    std::vector<int> pos(data.size());

    for (unsigned int i = 0; i < data.size(); i++) {
    
      seq_t map_seq(data[i]);
      trained_hmm.map(map_seq);
      //std::cout << "MAP sequence: " << std::endl;
      //map_seq.print();
  
      pos[i] = -1;
      for (unsigned int t = 0; t < map_seq.qseq.size(); t++) {
	if (map_seq.qseq[t] == 1 && map_seq.qseq[t+5] == 6) pos[i] = t;
      }
    }

    std::ofstream output(vm["output"].as<std::string>().c_str());
    for (unsigned int i = 0; i < data.size(); i++) {
      output << ">" << i+1 << " pos = " << pos[i] << std::endl;
      if (pos[i] != -1) {
	for (int k = 0; k < 20-pos[i]; k++) output << " ";
	for (int k = 0; k < pos[i]; k++) output << atcg[data[i][k]];
	output << " ";
	for (int k = pos[i]; k < pos[i]+5; k++) output << atcg[data[i][k]];
	output << " ";
	for (unsigned int k = pos[i]+5; k < data[i].size()-1; k++) output << atcg[data[i][k]];
	output << std::endl;
      }
      else {
	for (unsigned int k = 0; k < data[i].size(); k++) output << atcg[data[i][k]];
	output << std::endl;
      }
    }
    output.close();
  }
}

template<typename vseq_t>
void load_sequence(const char* fn, vseq_t& data, const char* atcg = "ATCG")
{
  std::ifstream input(fn);
  //int i = 1;
  while (!input.eof()) {
    std::string s;
    getline(input, s);
    //const char *atcg = "ATCG";
    if (s.length() > 0 && strchr(atcg, s[0])) {
      //std::cout << i++ << " " << s  << std::endl;
      typename vseq_t::value_type seq(s.length()+1);
      for (unsigned int i = 0; i < s.length(); i++) {
	const char* pos = strchr(atcg, s[i]);
	if (!pos) { std::cerr << "error unknown char " << s[i] << std::endl; exit(1); }
	seq[i] = pos - atcg;
      }
      seq[s.length()] = 4;
      data.push_back(seq);
    }
  }
}

void real_motif_test()
{
  std::cout << " ================ Real Motif with Duration HMM ================ " << std::endl;
  typedef DurationHMM hmm_t;
  typedef hmm_t::joint_seq_t seq_t;
  typedef hmm_t::vseq_t vseq_t;
  hmm_t hmm;
  hmm.load(vm["hmmfile"].as<std::string>().c_str());

  // int N = 1000;
  // int T = 8;
  vseq_t data;
  
  load_sequence(vm["seqfile"].as<std::string>().c_str(), data);

  hmm_t trained_hmm = hmm;
  srand(time(NULL));
  trained_hmm.rand();
  std::cout << "Before: " << std::endl;
  trained_hmm.trans_ = hmm.trans_; // fix the transmission probability table
  trained_hmm.start_ = hmm.start_; // and the starting probabilities
  trained_hmm.dprob_ = hmm.dprob_;
  for (int j = 0; j < trained_hmm.n_state(); j++) {
    trained_hmm.emis_(trained_hmm.n_observation()-1, j) = 0.0;
    if (j == trained_hmm.n_state()-1) {
      trained_hmm.emis_(trained_hmm.n_observation()-1, j) = 1.0;
      for (int k = 0; k < trained_hmm.n_observation()-1; k++)
	trained_hmm.emis_(k, j) = 0.0;
    }
    trained_hmm.emis_.col(j) /= arma::sum(trained_hmm.emis_.col(j));
  }
  trained_hmm.print_params();

  trained_hmm.set_prefix(vm["prefix"].as<std::string>());
  trained_hmm.train(data, vm["iter"].as<int>());
  std::cout << "Trained HMM: " << std::endl;
  trained_hmm.print_params();

  std::vector<int> pos(data.size());
  for (unsigned int i = 0; i < data.size(); i++) {
    
    seq_t map_seq(data[i]);
    trained_hmm.map(map_seq);
    //std::cout << "MAP sequence: " << std::endl;
    //map_seq.print();
  
    pos[i] = -1;
    for (unsigned int t = 0; t < map_seq.qseq.size(); t++) {
      if (map_seq.qseq[t] == 1 && map_seq.qseq[t+5] == 6) pos[i] = t;
    }
  }

  std::ofstream output(vm["output"].as<std::string>().c_str());
  const char *atcg = "ATCG";
  for (unsigned int i = 0; i < data.size(); i++) {
    output << ">" << i+1 << " pos = " << pos[i] << std::endl;
    if (pos[i] != -1) {
      for (int k = 0; k < 50-pos[i]; k++) output << " ";
      for (int k = 0; k < pos[i]; k++) output << atcg[data[i][k]];
      output << " ";
      for (int k = pos[i]; k < pos[i]+6; k++) output << atcg[data[i][k]];
      output << " ";
      for (unsigned int k = pos[i]+6; k < data[i].size()-1; k++) output << atcg[data[i][k]];
      output << std::endl;
    }
    else {
      for (unsigned int k = 0; k < data[i].size(); k++) output << atcg[data[i][k]];
      output << std::endl;
    }
  }
  output.close();
}

void position_test()
{
  std::cout << " ================ Real Motif with Duration HMM ================ " << std::endl;
  typedef PositionModel pm_t;
  typedef pm_t::vseq_t vseq_t;

  vseq_t data;
  load_sequence(vm["seqfile"].as<std::string>().c_str(), data);
  for (unsigned int i = 0; i < data.size(); i++) data[i].pop_back();

  srand(time(NULL));
  int motif_length = 6;
  int n_observation = 4;
  int max_seq_length = 50;
  pm_t pm(motif_length, n_observation, max_seq_length);
  pm.print_params();

  pm.train(data, vm["iter"].as<int>());
  std::cout << " -- Trained --" << std::endl;
  pm.print_params();

  std::vector<int> pos(data.size());
  for (unsigned int i = 0; i < data.size(); i++)
    pos[i] = pm.map_pos(data[i]);

  std::ofstream output(vm["output"].as<std::string>().c_str());
  const char *atcg = "ATCG";
  for (unsigned int i = 0; i < data.size(); i++) {
    output << ">" << i+1 << " pos = " << pos[i] << std::endl;
    if (pos[i] != -1) {
      for (int k = 0; k < 50-pos[i]; k++) output << " ";
      for (int k = 0; k < pos[i]; k++) output << atcg[data[i][k]];
      output << " ";
      for (int k = pos[i]; k < pos[i]+motif_length; k++) output << atcg[data[i][k]];
      output << " ";
      for (unsigned int k = pos[i]+motif_length; k < data[i].size(); k++) output << atcg[data[i][k]];
      output << std::endl;
    }
    else {
      for (unsigned int k = 0; k < data[i].size(); k++) output << atcg[data[i][k]];
      output << std::endl;
    }
  }
  output.close();
  
}

int main(int argc, char** argv)
{
  std::cout << "hmm tests ..." << std::endl;

  process_options(argc, argv);

  boost::posix_time::ptime time_start(boost::posix_time::microsec_clock::local_time());

  if (vm.count("help"))
    usage();
  else
  {
    if (vm.count("classic"))
      classic_hmm_test();
    else if (vm.count("duration"))
      duration_hmm_test();
    else if (vm.count("compare"))
      duration_hmm_compare_test();
    else if (vm.count("duration_markov"))
      duration_hmm_markov_test();
    else if (vm.count("compare_markov"))
      duration_hmm_markov_compare_test();
    else if (vm.count("motif"))
      motif_test();
    else if (vm.count("real"))
      real_motif_test();
    else if (vm.count("position"))
      position_test();
    else if (vm.count("anneal"))
      anneal_test();
  }

  boost::posix_time::ptime time_end(boost::posix_time::microsec_clock::local_time());
  boost::posix_time::time_duration duration(time_end - time_start);
  std::cerr << "Duration = " << duration << "\n";
}

void process_options(int argc, char** argv)
{
  desc.add_options()
    ("help", "produce this usage")
    ("hmmfile", po::value<std::string>()->default_value("hmm1.txt"), "input file (hmm)")
    ("seqfile", po::value<std::string>(), "sequence file")
    ("output", po::value<std::string>()->default_value("output.txt"), "output file")
    ("prefix", po::value<std::string>()->default_value("output/"), "output prefix")
    ("iter", po::value<int>()->default_value(10), "number of iterations")    
//("fin", po::value<std::string>(&fin)->default_value(""), "input file (factor functions)")
    //("out", po::value<std::string>(&out)->default_value(""), "output file")
    ("classic", "test the classic discrete HMM")
    ("compare", "convergence rate")
    ("duration", "test the duration discrete HMM (independent)")
    ("compare_markov", "convergence rate (markov)")
    ("duration_markov", "test the duration discrete HMM (markov)")
    ("motif", "motif finding")
    ("real", "real motif finding")
    ("position", "position model for motif finding")
    ("anneal", "simulated annealing for motif finding")
    ;
  po::store(po::parse_command_line(argc, argv, desc), vm);
  po::notify(vm);
}

